{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi-Agent Reinforcement Learning (MARL)\n",
    "\n",
    "This tutorial demonstrates how to use Tianshou for multi-agent reinforcement learning scenarios. We'll explore different MARL paradigms and implement a practical example using the Tic-Tac-Toe game."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MARL Paradigms\n",
    "\n",
    "Tianshou supports three fundamental types of multi-agent reinforcement learning paradigms:\n",
    "\n",
    "1. **Simultaneous move**: All agents take their actions at each timestep simultaneously (e.g., MOBA games)\n",
    "2. **Cyclic move**: Agents take actions sequentially in turns (e.g., Go)\n",
    "3. **Conditional move**: The environment conditionally selects which agent acts at each timestep (e.g., [Pig Game](https://en.wikipedia.org/wiki/Pig_(dice_game)))\n",
    "\n",
    "Our approach addresses these multi-agent RL problems by converting them into traditional single-agent RL formulations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Converting MARL to Single-Agent RL\n",
    "\n",
    "### Simultaneous Move\n",
    "\n",
    "For simultaneous-move scenarios, the solution is straightforward: we add an extra `num_agents` dimension to the state, action, and reward tensors. No other modifications are necessary.\n",
    "\n",
    "### Cyclic and Conditional Move\n",
    "\n",
    "Both cyclic and conditional move scenarios can be unified into a single framework. At each timestep, the environment selects an agent identified by `agent_id` to act. Since multiple agents are typically wrapped into a single object (the \"abstract agent\"), we pass the `agent_id` to this abstract agent, which then delegates the action to the appropriate specific agent.\n",
    "\n",
    "Additionally, in multi-agent RL, the set of legal actions often varies across timesteps (as in Go). Therefore, the environment must also provide a legal action mask to the abstract agent. This mask is a boolean array where `True` indicates available actions and `False` indicates illegal actions at the current timestep.\n",
    "\n",
    "<div style=\"text-align: center; padding: 1rem;\">\n",
    "<img src=\"../_static/images/marl.png\" style=\"height: 300px; padding-bottom: 1rem;\"><br>\n",
    "The abstract agent framework for multi-agent RL\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Unified Formulation\n",
    "\n",
    "This architecture leads to the following formulation of multi-agent RL:\n",
    "\n",
    "```python\n",
    "act = policy(state, agent_id, mask)\n",
    "(next_state, next_agent_id, next_mask), reward = env.step(act)\n",
    "```\n",
    "\n",
    "By constructing an augmented state `state_ = (state, agent_id, mask)`, we can reduce this to the standard single-agent RL formulation:\n",
    "\n",
    "```python\n",
    "act = policy(state_)\n",
    "next_state_, reward = env.step(act)\n",
    "```\n",
    "\n",
    "Following this principle, we'll implement a Q-learning algorithm to play [Tic-Tac-Toe](https://en.wikipedia.org/wiki/Tic-tac-toe) against a random opponent."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PettingZoo Integration\n",
    "\n",
    "Tianshou is fully compatible with [PettingZoo](https://pettingzoo.farama.org/) environments for multi-agent RL. While Tianshou doesn't directly provide specialized MARL facilities, it offers a flexible framework that can be adapted to various MARL scenarios.\n",
    "\n",
    "For comprehensive tutorials on using Tianshou with PettingZoo, refer to:\n",
    "\n",
    "* [Beginner Tutorial](https://pettingzoo.farama.org/tutorials/tianshou/beginner/)\n",
    "* [Intermediate Tutorial](https://pettingzoo.farama.org/tutorials/tianshou/intermediate/)\n",
    "* [Advanced Tutorial](https://pettingzoo.farama.org/tutorials/tianshou/advanced/)\n",
    "\n",
    "In this tutorial, we'll demonstrate how to use Tianshou in a multi-agent setting where only one agent is trained while the other uses a fixed random policy. You can then use this as a blueprint to replace the random policy with another trainable agent.\n",
    "\n",
    "Specifically, we'll train an agent to play Tic-Tac-Toe against a random opponent:\n",
    "\n",
    "<div style=\"text-align: center; padding: 1rem;\">\n",
    "<img src=\"../_static/images/tic-tac-toe.png\" style=\"padding-bottom: 1rem;\"><br>\n",
    "Tic-Tac-Toe game board\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exploring the Tic-Tac-Toe Environment\n",
    "\n",
    "The complete scripts are located in `test/pettingzoo/`. Tianshou provides the `PettingZooEnv` wrapper class that can wrap any PettingZoo environment. Let's explore the 3×3 Tic-Tac-Toe environment provided by PettingZoo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pettingzoo.classic import tictactoe_v3  # the Tic-Tac-Toe environment\n",
    "\n",
    "from tianshou.env import PettingZooEnv  # wrapper for PettingZoo environments\n",
    "\n",
    "# Initialize the environment\n",
    "# The board has 3 rows and 3 columns (9 positions total)\n",
    "# Players place 'X' and 'O' alternately on the board\n",
    "# The first player to get 3 consecutive marks wins\n",
    "env = PettingZooEnv(tictactoe_v3.env(render_mode=\"human\"))\n",
    "obs = env.reset()\n",
    "env.render()  # render the empty board"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output shows an empty 3×3 board:\n",
    "\n",
    "```\n",
    "board (step 0):\n",
    "     |     |\n",
    "  -  |  -  |  -\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  -  |  -  |  -\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  -  |  -  |  -\n",
    "     |     |\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Examine the observation structure\n",
    "print(obs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Understanding the Observation Space\n",
    "\n",
    "The observation returned by the environment is a dictionary with three keys:\n",
    "\n",
    "- **`agent_id`**: The identifier of the currently acting agent (e.g., `'player_1'` or `'player_2'`)\n",
    "\n",
    "- **`obs`**: The actual environment observation. For Tic-Tac-Toe, this is a numpy array with shape `(3, 3, 2)`:\n",
    "  - For `player_1`: The first 3×3 plane represents X placements, the second plane represents O placements\n",
    "  - For `player_2`: The planes are swapped (O in first plane, X in second)\n",
    "  - Each cell contains either 0 (empty/not placed) or 1 (mark placed)\n",
    "\n",
    "- **`mask`**: A boolean array indicating legal actions at the current timestep. For Tic-Tac-Toe, index `i` corresponds to position `(i // 3, i % 3)` on the board. If `mask[i] == True`, the player can place their mark at that position. Initially, all positions are available, so all mask values are `True`.\n",
    "\n",
    "> **Note**: The mask representation is flexible and works for both discrete and continuous action spaces. While we use a boolean array here, you could also use action spaces like `gymnasium.spaces.Discrete` or `gymnasium.spaces.Box` to represent available actions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Playing a Few Steps\n",
    "\n",
    "Let's play a couple of moves to understand the environment dynamics better."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# Take an action (place mark at position 0 - top-left corner)\n",
    "action = 0  # action can be an integer or a numpy array with one element\n",
    "obs, reward, done, truncated, info = env.step(action)  # follows the Gymnasium API\n",
    "\n",
    "print(\"Observation after first move:\")\n",
    "print(obs)\n",
    "\n",
    "# Examine the reward structure\n",
    "# Reward has two items (one for each player): 1 for win, -1 for loss, 0 otherwise\n",
    "print(f\"\\nReward: {reward}\")\n",
    "\n",
    "# Check if the game is over\n",
    "print(f\"Done: {done}\")\n",
    "\n",
    "# Info is typically an empty dict in Tic-Tac-Toe but may contain useful information in other environments\n",
    "print(f\"Info: {info}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that after the first move:\n",
    "- The `agent_id` switches to `'player_2'`\n",
    "- The observation array shows the X placement in the first position\n",
    "- The mask now has `False` at index 0 (that position is occupied)\n",
    "- The reward is `[0, 0]` (no winner yet)\n",
    "- The game continues (`done = False`)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "Note: If we continue playing, the game terminates when only one empty position remains, rather than when the board is completely full. This is because a player with only one available position has no meaningful choice."
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Random Agents\n",
    "\n",
    "Now that we understand the environment, let's start by watching two random agents play against each other.\n",
    "\n",
    "Tianshou provides built-in classes for multi-agent learning. The key components are:\n",
    "\n",
    "- **`RandomPolicy`**: A policy that randomly selects actions\n",
    "- **`MultiAgentPolicyManager`**: Manages multiple agent policies and delegates actions to the appropriate agent based on `agent_id`\n",
    "\n",
    "<div style=\"text-align: center; padding: 1rem;\">\n",
    "<img src=\"../_static/images/marl.png\" style=\"height: 300px; padding-bottom: 1rem;\"><br>\n",
    "The relationship between MultiAgentPolicyManager and individual agent policies\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tianshou.algorithm.multiagent.marl import MultiAgentOffPolicyAlgorithm\n",
    "from tianshou.algorithm.random import MARLRandomDiscreteMaskedOffPolicyAlgorithm\n",
    "from tianshou.data import Collector\n",
    "from tianshou.env import DummyVectorEnv\n",
    "\n",
    "# Create a multi-agent algorithm with two random agents\n",
    "policy = MultiAgentOffPolicyAlgorithm(\n",
    "    algorithms=[\n",
    "        MARLRandomDiscreteMaskedOffPolicyAlgorithm(action_space=env.action_space),\n",
    "        MARLRandomDiscreteMaskedOffPolicyAlgorithm(action_space=env.action_space),\n",
    "    ],\n",
    "    env=env,\n",
    ")\n",
    "\n",
    "# Vectorize the environment for the collector\n",
    "env = DummyVectorEnv([lambda: env])\n",
    "\n",
    "# Create a collector to gather trajectories\n",
    "collector = Collector(policy, env)\n",
    "\n",
    "# Collect and visualize one episode\n",
    "result = collector.collect(n_episode=1, render=0.1, reset_before_collect=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You'll see the game progress step by step. Here's an example of the final moves:\n",
    "\n",
    "```\n",
    "     |     |\n",
    "  X  |  X  |  -\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  X  |  O  |  -\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  O  |  -  |  -\n",
    "     |     |\n",
    "```\n",
    "\n",
    "```\n",
    "     |     |\n",
    "  X  |  X  |  -\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  X  |  O  |  -\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  O  |  -  |  O\n",
    "     |     |\n",
    "```\n",
    "\n",
    "```\n",
    "     |     |\n",
    "  X  |  X  |  X\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  X  |  O  |  -\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  O  |  -  |  O\n",
    "     |     |\n",
    "```\n",
    "\n",
    "Random agents perform poorly. In the game above, although agent 2 eventually wins, a smart agent 1 would have won immediately by placing an X at position (1, 1) (center of middle row)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training an Agent Against a Random Opponent\n",
    "\n",
    "Now let's train an intelligent agent! We'll use Deep Q-Network (DQN) to learn optimal play against a random opponent."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports and Setup\n",
    "\n",
    "First, let's import all necessary modules:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from copy import deepcopy\n",
    "from functools import partial\n",
    "\n",
    "import gymnasium\n",
    "import torch\n",
    "from pettingzoo.classic import tictactoe_v3\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "\n",
    "from tianshou.algorithm import (\n",
    "    DQN,\n",
    "    Algorithm,\n",
    "    MARLRandomDiscreteMaskedOffPolicyAlgorithm,\n",
    "    MultiAgentOffPolicyAlgorithm,\n",
    ")\n",
    "from tianshou.algorithm.algorithm_base import OffPolicyAlgorithm\n",
    "from tianshou.algorithm.modelfree.dqn import DiscreteQLearningPolicy\n",
    "from tianshou.algorithm.optim import AdamOptimizerFactory, OptimizerFactory\n",
    "from tianshou.data import Collector, CollectStats, VectorReplayBuffer\n",
    "from tianshou.data.stats import InfoStats\n",
    "from tianshou.env import DummyVectorEnv\n",
    "from tianshou.env.pettingzoo_env import PettingZooEnv\n",
    "from tianshou.trainer import OffPolicyTrainerParams\n",
    "from tianshou.utils import TensorboardLogger\n",
    "from tianshou.utils.net.common import Net"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hyperparameters\n",
    "\n",
    "Let's define the hyperparameters for our training experiment directly (no argparse needed in notebooks!):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define hyperparameters\n",
    "class Args:\n",
    "    seed = 1626\n",
    "    eps_test = 0.05\n",
    "    eps_train = 0.1\n",
    "    buffer_size = 20000\n",
    "    lr = 1e-4\n",
    "    gamma = 0.9  # A smaller gamma favors earlier wins\n",
    "    n_step = 3\n",
    "    target_update_freq = 320\n",
    "    epoch = 50\n",
    "    epoch_num_steps = 1000\n",
    "    collection_step_num_env_steps = 10\n",
    "    update_per_step = 0.1\n",
    "    batch_size = 64\n",
    "    hidden_sizes = [128, 128, 128, 128]  # noqa: RUF012\n",
    "    num_train_envs = 10\n",
    "    num_test_envs = 10\n",
    "    logdir = \"log\"\n",
    "    render = 0.1\n",
    "    win_rate = 0.6  # Target winning rate (optimal policy can get ~0.7)\n",
    "    watch = False  # Set to True to skip training and watch pre-trained models\n",
    "    agent_id = 2  # The learned agent plays as player 2\n",
    "    resume_path = \"\"  # Path to pre-trained agent .pth file\n",
    "    opponent_path = \"\"  # Path to pre-trained opponent .pth file\n",
    "    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "    model_save_path = None  # Will be set in save_best_fn\n",
    "\n",
    "\n",
    "args = Args()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Agent Setup\n",
    "\n",
    "The `get_agents` function creates and configures our agents:\n",
    "\n",
    "- **Neural Network**: We use `Net`, a multi-layer perceptron with ReLU activations\n",
    "- **Learning Algorithm**: A `DiscreteQLearningPolicy` combined with `DQN` for Q-learning updates\n",
    "- **Opponent**: Either a `MARLRandomDiscreteMaskedOffPolicyAlgorithm` that randomly chooses legal actions, or a pre-trained agent for self-play\n",
    "\n",
    "Both agents are managed by `MultiAgentOffPolicyAlgorithm`, which:\n",
    "- Calls the correct agent based on `agent_id` in the observation\n",
    "- Dispatches data to each agent according to their `agent_id`\n",
    "- Makes each agent perceive the environment as a single-agent problem\n",
    "\n",
    "<div style=\"text-align: center; padding: 1rem;\">\n",
    "<img src=\"../_static/images/marl.png\" style=\"height: 300px; padding-bottom: 1rem;\"><br>\n",
    "How MultiAgentOffPolicyAlgorithm coordinates agent algorithms\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_env(render_mode: str | None = None) -> PettingZooEnv:\n",
    "    return PettingZooEnv(tictactoe_v3.env(render_mode=render_mode))\n",
    "\n",
    "\n",
    "def get_agents(\n",
    "    args,\n",
    "    agent_learn: OffPolicyAlgorithm | None = None,\n",
    "    agent_opponent: OffPolicyAlgorithm | None = None,\n",
    "    optim: OptimizerFactory | None = None,\n",
    ") -> tuple[MultiAgentOffPolicyAlgorithm, torch.optim.Optimizer | None, list]:\n",
    "    \"\"\"Create or load agents for training.\"\"\"\n",
    "    env = get_env()\n",
    "    observation_space = (\n",
    "        env.observation_space.spaces[\"observation\"]\n",
    "        if isinstance(env.observation_space, gymnasium.spaces.Dict)\n",
    "        else env.observation_space\n",
    "    )\n",
    "    args.state_shape = observation_space.shape or int(observation_space.n)\n",
    "    args.action_shape = env.action_space.shape or int(env.action_space.n)\n",
    "\n",
    "    if agent_learn is None:\n",
    "        # Create the neural network model\n",
    "        net = Net(\n",
    "            state_shape=args.state_shape,\n",
    "            action_shape=args.action_shape,\n",
    "            hidden_sizes=args.hidden_sizes,\n",
    "        ).to(args.device)\n",
    "\n",
    "        if optim is None:\n",
    "            optim = AdamOptimizerFactory(lr=args.lr)\n",
    "\n",
    "        # Create Q-learning policy for the learning agent\n",
    "        algorithm = DiscreteQLearningPolicy(\n",
    "            model=net,\n",
    "            action_space=env.action_space,\n",
    "            eps_training=args.eps_train,\n",
    "            eps_inference=args.eps_test,\n",
    "        )\n",
    "\n",
    "        # Wrap in DQN algorithm\n",
    "        agent_learn = DQN(\n",
    "            policy=algorithm,\n",
    "            optim=optim,\n",
    "            n_step_return_horizon=args.n_step,\n",
    "            gamma=args.gamma,\n",
    "            target_update_freq=args.target_update_freq,\n",
    "        )\n",
    "\n",
    "        if args.resume_path:\n",
    "            agent_learn.load_state_dict(torch.load(args.resume_path))\n",
    "\n",
    "    if agent_opponent is None:\n",
    "        if args.opponent_path:\n",
    "            # Load a pre-trained opponent for self-play\n",
    "            agent_opponent = deepcopy(agent_learn)\n",
    "            agent_opponent.load_state_dict(torch.load(args.opponent_path))\n",
    "        else:\n",
    "            # Use a random opponent\n",
    "            agent_opponent = MARLRandomDiscreteMaskedOffPolicyAlgorithm(\n",
    "                action_space=env.action_space\n",
    "            )\n",
    "\n",
    "    # Arrange agents based on which player position the learning agent takes\n",
    "    if args.agent_id == 1:\n",
    "        agents = [agent_learn, agent_opponent]\n",
    "    else:\n",
    "        agents = [agent_opponent, agent_learn]\n",
    "\n",
    "    ma_algorithm = MultiAgentOffPolicyAlgorithm(algorithms=agents, env=env)\n",
    "    return ma_algorithm, optim, env.agents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training Loop\n",
    "\n",
    "The training procedure follows the standard Tianshou workflow, similar to single-agent DQN training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_agent(\n",
    "    args,\n",
    "    agent_learn: OffPolicyAlgorithm | None = None,\n",
    "    agent_opponent: OffPolicyAlgorithm | None = None,\n",
    "    optim: OptimizerFactory | None = None,\n",
    ") -> tuple[InfoStats, OffPolicyAlgorithm]:\n",
    "    \"\"\"Train the agent using DQN.\"\"\"\n",
    "    # ======== Environment Setup =========\n",
    "    train_envs = DummyVectorEnv([get_env for _ in range(args.num_train_envs)])\n",
    "    test_envs = DummyVectorEnv([get_env for _ in range(args.num_test_envs)])\n",
    "\n",
    "    # Set random seeds for reproducibility\n",
    "    np.random.seed(args.seed)\n",
    "    torch.manual_seed(args.seed)\n",
    "    train_envs.seed(args.seed)\n",
    "    test_envs.seed(args.seed)\n",
    "\n",
    "    # ======== Agent Setup =========\n",
    "    marl_algorithm, optim, agents = get_agents(\n",
    "        args,\n",
    "        agent_learn=agent_learn,\n",
    "        agent_opponent=agent_opponent,\n",
    "        optim=optim,\n",
    "    )\n",
    "\n",
    "    # ======== Collector Setup =========\n",
    "    train_collector = Collector[CollectStats](\n",
    "        marl_algorithm,\n",
    "        train_envs,\n",
    "        VectorReplayBuffer(args.buffer_size, len(train_envs)),\n",
    "        exploration_noise=True,\n",
    "    )\n",
    "    test_collector = Collector[CollectStats](marl_algorithm, test_envs, exploration_noise=True)\n",
    "\n",
    "    # Collect initial random samples\n",
    "    train_collector.reset()\n",
    "    train_collector.collect(n_step=args.batch_size * args.num_train_envs)\n",
    "\n",
    "    # ======== Logging Setup =========\n",
    "    log_path = os.path.join(args.logdir, \"tic_tac_toe\", \"dqn\")\n",
    "    writer = SummaryWriter(log_path)\n",
    "    writer.add_text(\"args\", str(args))\n",
    "    logger = TensorboardLogger(writer)\n",
    "\n",
    "    player_agent_id = agents[args.agent_id - 1]\n",
    "\n",
    "    # ======== Callback Functions =========\n",
    "    def save_best_fn(policy: Algorithm) -> None:\n",
    "        \"\"\"Save the best performing policy.\"\"\"\n",
    "        if hasattr(args, \"model_save_path\") and args.model_save_path:\n",
    "            model_save_path = args.model_save_path\n",
    "        else:\n",
    "            model_save_path = os.path.join(args.logdir, \"tic_tac_toe\", \"dqn\", \"policy.pth\")\n",
    "        torch.save(policy.get_algorithm(player_agent_id).state_dict(), model_save_path)\n",
    "\n",
    "    def stop_fn(mean_rewards: float) -> bool:\n",
    "        \"\"\"Stop training when target win rate is achieved.\"\"\"\n",
    "        return mean_rewards >= args.win_rate\n",
    "\n",
    "    def reward_metric(rews: np.ndarray) -> np.ndarray:\n",
    "        \"\"\"Extract the reward for our learning agent.\"\"\"\n",
    "        return rews[:, args.agent_id - 1]\n",
    "\n",
    "    # ======== Trainer =========\n",
    "    result = marl_algorithm.run_training(\n",
    "        OffPolicyTrainerParams(\n",
    "            train_collector=train_collector,\n",
    "            test_collector=test_collector,\n",
    "            max_epochs=args.epoch,\n",
    "            epoch_num_steps=args.epoch_num_steps,\n",
    "            collection_step_num_env_steps=args.collection_step_num_env_steps,\n",
    "            test_step_num_episodes=args.num_test_envs,\n",
    "            batch_size=args.batch_size,\n",
    "            stop_fn=stop_fn,\n",
    "            save_best_fn=save_best_fn,\n",
    "            update_step_num_gradient_steps_per_sample=args.update_per_step,\n",
    "            logger=logger,\n",
    "            test_in_train=False,\n",
    "            multi_agent_return_reduction=reward_metric,\n",
    "            show_progress=False,\n",
    "        )\n",
    "    )\n",
    "\n",
    "    return result, marl_algorithm.get_algorithm(player_agent_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation Function\n",
    "\n",
    "This function allows us to watch a trained agent play:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def watch(\n",
    "    args,\n",
    "    agent_learn: OffPolicyAlgorithm | None = None,\n",
    "    agent_opponent: OffPolicyAlgorithm | None = None,\n",
    ") -> None:\n",
    "    \"\"\"Watch a pre-trained agent play.\"\"\"\n",
    "    env = DummyVectorEnv([partial(get_env, render_mode=\"human\")])\n",
    "    policy, optim, agents = get_agents(args, agent_learn=agent_learn, agent_opponent=agent_opponent)\n",
    "    collector = Collector[CollectStats](policy, env, exploration_noise=True)\n",
    "    result = collector.collect(n_episode=1, render=args.render, reset_before_collect=True)\n",
    "    result.pprint_asdict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running the Training\n",
    "\n",
    "Now let's train the agent and watch it play!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train the agent\n",
    "result, agent = train_agent(args)\n",
    "\n",
    "# Watch the trained agent play\n",
    "watch(args, agent)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Results\n",
    "\n",
    "After training for less than a minute, you'll see the agent play against the random opponent. Here's an example game:\n",
    "\n",
    "<details>\n",
    "<summary>Example: Trained Agent vs Random Opponent</summary>\n",
    "\n",
    "```\n",
    "     |     |\n",
    "  -  |  -  |  -\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  -  |  -  |  X\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  -  |  -  |  -\n",
    "     |     |\n",
    "```\n",
    "\n",
    "```\n",
    "     |     |\n",
    "  -  |  -  |  -\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  -  |  O  |  X\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  -  |  -  |  -\n",
    "     |     |\n",
    "```\n",
    "\n",
    "```\n",
    "     |     |\n",
    "  -  |  -  |  -\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  X  |  O  |  X\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  -  |  -  |  -\n",
    "     |     |\n",
    "```\n",
    "\n",
    "```\n",
    "     |     |\n",
    "  -  |  O  |  -\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  X  |  O  |  X\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  -  |  -  |  -\n",
    "     |     |\n",
    "```\n",
    "\n",
    "```\n",
    "     |     |\n",
    "  -  |  O  |  -\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  X  |  O  |  X\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  -  |  X  |  -\n",
    "     |     |\n",
    "```\n",
    "\n",
    "```\n",
    "     |     |\n",
    "  O  |  O  |  -\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  X  |  O  |  X\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  -  |  X  |  -\n",
    "     |     |\n",
    "```\n",
    "\n",
    "```\n",
    "     |     |\n",
    "  O  |  O  |  X\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  X  |  O  |  X\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  -  |  X  |  -\n",
    "     |     |\n",
    "```\n",
    "\n",
    "```\n",
    "     |     |\n",
    "  O  |  O  |  X\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  X  |  O  |  X\n",
    "_____|_____|_____\n",
    "     |     |\n",
    "  -  |  X  |  O\n",
    "     |     |\n",
    "```\n",
    "\n",
    "Final reward: 1.0, length: 8.0\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that our trained agent plays as player 2 (O) and wins! The agent has learned the game rules through trial and error, understanding that three consecutive O marks lead to victory."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "It is easily possible to make the trained agent play against itself. Try this as an exercise!"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While the trained agent plays well against a random opponent, it's still far from perfect play. The next step would be to implement self-play training, similar to AlphaZero, where the agent continuously improves by playing against increasingly stronger versions of itself."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In this tutorial, we demonstrated how to use Tianshou for training a single agent in a multi-agent reinforcement learning setting. Key takeaways:\n",
    "\n",
    "1. **MARL Paradigms**: Tianshou supports simultaneous, cyclic, and conditional move scenarios\n",
    "2. **Abstraction**: Multi-agent problems can be converted to single-agent RL through clever state augmentation\n",
    "3. **PettingZoo Integration**: Seamless compatibility with PettingZoo environments via `PettingZooEnv`\n",
    "4. **Algorithm Management**: `MultiAgentOffPolicyAlgorithm` handles agent coordination and data distribution\n",
    "5. **Flexible Framework**: Easy to extend from single-agent training to more complex multi-agent scenarios\n",
    "\n",
    "Tianshou provides a flexible and intuitive framework for reinforcement learning. Experiment with different architectures, training regimes, and opponent strategies to build even more capable agents!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
